import cvxpy as cp
import gurobipy as gp
import numpy as np

from src.util.util import *


class ADMM:

    def __init__(self, env):
        self.env = env
        # u = {a, B, f}, v = {x, y}
        action = np.ones(env.M, dtype=int)
        bandwidth, frequency = np.ones(env.M), np.ones(env.M)
        for ues_idx in env.parts_ues:
            bandwidth[ues_idx] = env.config.uav_bandwidth / len(ues_idx)
            frequency[ues_idx] = env.config.uav_cpu_frequency / len(ues_idx)
        self.u = np.stack((action, bandwidth, frequency), axis=1)
        self.v = np.stack((bandwidth.copy(), frequency.copy()), axis=1)
        self.alpha, self.beta, = -np.zeros(env.M), -np.zeros(env.M)
        self.c = -0.001

    def run(self, collector):
        for i in range(1500):
            print("**************************************************************************")
            print(f"Current episode: {i}")
            # 1. 联合优化
            self.run_epoch(epoch=i)
            # 2. 计算当前决策的总成本
            self.env.execute(action=self.u[:, 0], bandwidth=self.u[:, 1], frequency=self.u[:, 2])
            cost = self.env.calc_cost()
            # 3. 收集数据（成本 & 资源利用率）
            collector.collect(cost, int(sum(self.u[:, 0])))
            collector.next()
            print(f"Offloading cost: {cost}, Offloading num: {sum(self.u[:, 0])}, Alpha: {sum(self.alpha)}, Beta: {sum(self.beta)}")
            ratio_b = 0
            ratio_f = 0
            for g in self.env.parts_ues:
                ratio_b += sum(self.u[g, 1]) / self.env.config.uav_bandwidth
                ratio_f += sum(self.u[g, 2]) / self.env.config.uav_cpu_frequency
            print(f"Avg resource 1: {ratio_b / self.env.G}, 2: {ratio_f / self.env.G}")
            # 4. 停止检查
            err_b = self.u[:, 1] - self.v[:, 0]
            err_f = self.u[:, 2] - self.v[:, 1]
            err_b = [abs(x) for x in err_b if abs(x) > 1e-10]
            err_f = [abs(x) for x in err_f if abs(x) > 1e-10]
            print(f"Avd error, b: {sum(err_b)/len(err_b) if len(err_b) != 0 else 0}, f: {sum(err_f)/len(err_f) if len(err_f) != 0 else 0}")
        return self.u[:, 0], self.u[:, 1], self.u[:, 2]

    def f_b(self, idx, b):
        return ((self.env.meta_trans_rate[idx] / b) + (self.alpha[idx] * b) - ((self.c / 2) * ((b - self.v[(idx, 0)]) ** 2))) / self.env.meta_ue_loc[idx]

    def get_solution_b(self, solutions, idx, lb=.001, ub=1):
        solutions.append(lb)
        solutions.append(ub)
        v_solutions = [[s, self.f_b(idx, s)] for s in solutions if lb <= s <= ub]
        v_solutions = np.array(v_solutions)
        return v_solutions[v_solutions[:, 1].argsort()][(0, 0)]

    def f_f(self, idx, f):
        return ((self.env.meta_task_cpu[idx] / f) + (self.beta[idx] * f) - ((self.c / 2) * ((f - self.v[(idx, 1)]) ** 2))) / self.env.meta_ue_loc[idx]

    def get_solution_f(self, solutions, idx, lb=.001, ub=1):
        solutions.append(lb)
        solutions.append(ub)
        v_solutions = [[s, self.f_f(idx, s)] for s in solutions if lb <= s <= ub]
        v_solutions = np.array(v_solutions)
        return v_solutions[v_solutions[:, 1].argsort()][(0, 0)]

    def run_epoch(self, epoch):
        # 1. Given {v, w}, Calculate {u'}
        u = self.u
        for i in range(self.env.M):
            # 1.1. Bandwidth
            b_b = -(self.c * self.v[(i, 0)] + self.alpha[i]) / self.c
            b_d = (self.env.meta_trans_rate[i] / self.env.meta_ue_loc[i]) / self.c
            b_solutions = sheng_jin_formula(a=1, b=b_b, c=0, d=b_d)
            b = self.get_solution_b(b_solutions, idx=i, lb=.001, ub=self.env.config.uav_bandwidth)
            # 1.2. Frequency
            f_b = -(self.c * self.v[(i, 1)] + self.beta[i]) / self.c
            f_d = (self.env.meta_task_cpu[i] / self.env.meta_ue_loc[i]) / self.c
            f_solutions = sheng_jin_formula(a=1, b=f_b, c=0, d=f_d)
            f = self.get_solution_f(f_solutions, idx=i, lb=.001, ub=self.env.config.uav_cpu_frequency)
            # 1.3. Select local or edge
            cost_1 = self.env.parts[self.env.ues[i].part].t_wait + (self.env.meta_trans_rate[i] / b) + (self.env.meta_task_cpu[i] / f)
            cost_1 = cost_1 / self.env.meta_ue_loc[i]
            cost_2 = self.alpha[i] * (b - self.v[(i, 0)]) + self.beta[i] * (f - self.v[(i, 1)])
            cost_3 = (-self.c / 2) * (((b - self.v[(i, 0)]) ** 2) + ((f - self.v[(i, 1)]) ** 2))
            cost = cost_1 + cost_2 + cost_3
            cost_loc = 1 + cost_2 + (-self.c / 2) * ((self.v[(i, 0)] ** 2) + (self.v[(i, 1)] ** 2))
            u[i] = [0, 0, 0] if cost > cost_loc else [1, b, f]
        # 2. Given {u', w}, Calculate {v'}
        v = self.cvx_v(u)
        # 3. Given {u', v'}, Calculate {w'}
        alpha = self.alpha - self.c * (u[:, 1] - v[:, 0])
        beta = self.beta - self.c * (u[:, 2] - v[:, 1])
        # 4. Save epoch result
        self.u, self.v, self.alpha, self.beta = u, v, alpha, beta

    def cvx_u(self, ue_idx):
        cvx_model = gp.Model()
        b = cvx_model.addVar(lb=0.001, ub=self.env.config.uav_bandwidth, vtype=gp.GRB.CONTINUOUS, name=f"b@{ue_idx}")
        b_ = cvx_model.addVar(lb=1/self.env.config.uav_bandwidth, ub=1000, vtype=gp.GRB.CONTINUOUS, name=f"b_@{ue_idx}")
        f = cvx_model.addVar(lb=0.001, ub=self.env.config.uav_cpu_frequency[0], vtype=gp.GRB.CONTINUOUS, name=f"f@{ue_idx}")
        f_ = cvx_model.addVar(lb=1/self.env.config.uav_cpu_frequency[0], ub=1000, vtype=gp.GRB.CONTINUOUS, name=f"f_@{ue_idx}")
        cvx_model.addConstr(b * b_ == 1, name=f"bb_@{ue_idx}")
        cvx_model.addConstr(f * f_ == 1, name=f"ff_@{ue_idx}")
        obj = (1 / self.env.meta_ue_loc[ue_idx]) * (self.env.parts[self.env.ues[ue_idx].part].t_wait.sum() + (self.env.meta_trans_rate[ue_idx] * b_) + (self.env.meta_task_cpu[ue_idx] * f_))
        obj += (self.alpha[ue_idx] * b + self.beta[ue_idx] * f)
        obj += (-self.c / 2) * (((b - self.v[(ue_idx, 0)]) ** 2) + ((f - self.v[(ue_idx, 1)]) ** 2))
        cvx_model.setObjective(obj, sense=gp.GRB.MINIMIZE)
        cvx_model.Params.OutputFlag = 0
        cvx_model.Params.NonConvex = 2
        cvx_model.optimize()
        if cvx_model.status == gp.GRB.OPTIMAL:
            return b.X, f.X
        return 0

    def cvx_v(self, u):
        x_opt = np.zeros(self.env.M, dtype=int)
        y_opt = np.zeros(self.env.M, dtype=int)
        for g in range(self.env.G):
            # 1. 获取 PT 信息
            ues_idx_part = self.env.parts_ues[g]
            ues_idx_loc = [x for x in ues_idx_part if not int(u[(x, 0)])]
            ues_idx = [x for x in ues_idx_part if int(u[(x, 0)])]
            ues_num = len(ues_idx)
            # 2. 创建决策变量 v = {x, y}
            cvx_model = gp.Model()
            x = cvx_model.addVars(ues_num, lb=0, ub=self.env.config.uav_bandwidth, vtype=gp.GRB.CONTINUOUS, name="x")
            y = cvx_model.addVars(ues_num, lb=0, ub=self.env.config.uav_cpu_frequency, vtype=gp.GRB.CONTINUOUS, name="y")
            x_ = cvx_model.addVars(ues_num, lb=1/self.env.config.uav_bandwidth, vtype=gp.GRB.CONTINUOUS, name="x_")
            y_ = cvx_model.addVars(ues_num, lb=1/self.env.config.uav_cpu_frequency, vtype=gp.GRB.CONTINUOUS, name="y_")
            # 3. 设置约束条件
            # cvx_model.addConstr(gp.quicksum(x) <= self.env.config.uav_bandwidth, name="x_sum")
            # cvx_model.addConstr(gp.quicksum(y) <= self.env.config.uav_cpu_frequency, name="y_sum")
            for i in range(ues_num):
                # noinspection PyTypeChecker
                cvx_model.addConstr(x[i] * x_[i] == 1, name=f"zx_x@{i}")
                # noinspection PyTypeChecker
                cvx_model.addConstr(y[i] * y_[i] == 1, name=f"zy_y@{i}")
                c1 = self.env.meta_trans_rate[ues_idx[i]] * x_[i]
                c2 = self.env.meta_task_cpu[ues_idx[i]] * y_[i]
                cvx_model.addConstr(c1 + c2 <= self.env.config.uav_hover_time, name=f"hover@{i}")
            # 4. 创建优化问题
            obj = gp.quicksum(self.alpha[ues_idx[i]] * (u[(ues_idx[i], 1)] - x[i]) + self.beta[ues_idx[i]] * (u[(ues_idx[i], 2)] - y[i]) - (self.c/2) * (((u[(ues_idx[i], 1)] - x[i]) ** 2) + ((u[(ues_idx[i], 2)] - y[i]) ** 2)) for i in range(ues_num))
            cvx_model.setObjective(obj, sense=gp.GRB.MINIMIZE)
            cvx_model.Params.OutputFlag = 0
            cvx_model.Params.NonConvex = 2
            cvx_model.optimize()
            # 5. 获取结果
            if cvx_model.status == gp.GRB.OPTIMAL:
                for i in range(ues_num):
                    x_opt[ues_idx[i]] = x[i].X
                    y_opt[ues_idx[i]] = y[i].X
        return np.stack((x_opt, y_opt), axis=1)
